package elastic

import (
	"chatgpt/helper"
	"chatgpt/model/glm/request"
	"context"
	"github.com/olivere/elastic/v7"
	"github.com/slclub/go-tips/json"
	"net/http"
	"reflect"
	"time"
)

type SearchClient struct {
	Client *elastic.Client
	Index  string
	Maping string
}

func NewSearch(name string, mapping string) *SearchClient {
	client := &SearchClient{
		Index:  name,
		Maping: mapping,
	}
	client.Init()
	return client
}

func NewSearchFile(name string, mapping string) *SearchClient {
	client := &SearchClient{
		Index:  name,
		Maping: mapping,
	}
	client.Init()
	return client
}

func (self *SearchClient) Init() {
	cnf := helper.Config().Sub("elastic")
	hosts := cnf.GetStringSlice("hosts")
	client, err := elastic.NewClient(
		elastic.SetURL(hosts...),
		elastic.SetSniff(false),
		elastic.SetHealthcheckInterval(10*time.Second),
		//elastic.SetRetrier(),
		elastic.SetGzip(true),
		elastic.SetErrorLog(&logt{}),
		elastic.SetInfoLog(new(logt)),
		elastic.SetBasicAuth(cnf.GetString("account"), cnf.GetString("password")),
		elastic.SetHeaders(http.Header{
			"X-Caller-Id": []string{"..."},
		}),
	)
	if err != nil {
		return
	}
	self.Client = client
}

func (self *SearchClient) Search(q string, field string) []map[string]any {
	ctx := context.Background()
	matchQuery := elastic.NewMatchQuery(field, q).Analyzer("ik_max_word").Operator("or")
	searchResult, err := self.Client.Search().
		Index(self.Index).
		Query(matchQuery).
		//Source(matchQuery).
		Do(ctx)
	if err != nil {
		// Handle error
		logtt.Print("ELASTIC.Search.MATCH", err)
		return nil
	}
	return self.ConvToArray(searchResult)
}

func (self *SearchClient) ConvToArray(result *elastic.SearchResult) []map[string]any {
	hitarr := result.Hits.Hits
	rtn := []map[string]any{}
	for _, v := range hitarr {
		vm := map[string]any{}
		err := json.Unmarshal(v.Source, &vm)
		if err != nil {
			continue
		}
		vm["Score"] = v.Score
		rtn = append(rtn, vm)
	}
	return rtn
}

func (self *SearchClient) Filter(data []map[string]any) []map[string]any {
	fate := false
	jump := 0
	for k, v := range data {
		if val, ok := v["Score"].(*float64); ok && *val >= 1 {
			fate = true
			continue
		}
		if fate {
			data[k-jump] = data[len(data)-1]
			data = data[:len(data)-1]
			jump++
		}
	}
	return data
}

func (self *SearchClient) Bigger(data []map[string]any) []map[string]any {
	jump := 0
	max := float64(0.0)
	for k, v := range data {
		val, ok := v["Score"].(*float64)
		if !ok {
			continue
		}

		if *val > 1 && max <= *val {
			max = *val
			continue
		}
		if b := max - *val; b > 0 && b < 1 {
			continue
		}
		if max == 0.0 && *val > 0 {
			continue
		}

		data[k-jump] = data[len(data)-1]
		data = data[:len(data)-1]
		jump++
	}
	return data
}

func (self *SearchClient) Clear() error {
	ctx := context.Background()
	// 删除索引
	_, err := self.Client.DeleteIndex(self.Index).Do(ctx)
	// Handle error
	if err != nil {
		if elasticError, ok := err.(*elastic.Error); !ok || elasticError.Status != 404 {
			logtt.Print(err)
			return err
		}
	}

	// 创建索引
	createIndex, err := self.Client.CreateIndex(self.Index).BodyString(self.Maping).Do(ctx)
	if err != nil {
		// Handle error
		logtt.Print(err)
		return err
	}
	if !createIndex.Acknowledged {
		// Not acknowledged
	}
	return nil
}

func (self *SearchClient) AddDocument(data any) (string, error) {
	ctx := context.Background()
	indexSrv := self.Client.Index().Index(self.Index)
	indexSrv = indexSrv.BodyJson(data)
	res, err := indexSrv.Do(ctx)
	if err != nil {
		logtt.Print(res, err)
		return "", err
	}
	return res.Id, nil
}

func (self *SearchClient) GetDocument(id string) map[string]any {
	ctx := context.Background()
	get1, err := self.Client.Get().
		Index(self.Index).
		//Type("tweet").
		Id("1").
		Do(ctx)
	if err != nil {
		logtt.Print(err)
		return nil
	}
	data := map[string]any{}
	json.Unmarshal(get1.Source, &data)
	return data
}

func (self *SearchClient) DeleteDocument(id string) error {
	ctx := context.Background()
	_, err := self.Client.Delete().
		Index(self.Index).
		Id(id).
		Do(ctx)
	if err != nil {
		return err
	}
	return nil
}

func (self *SearchClient) ListDocuments(offset, limit int) ([]map[string]any, int) {
	rtn := []map[string]any{}
	ctx := context.Background()
	searchResult, err := self.Client.Search().
		Index(self.Index).
		From(offset).Size(limit).
		Do(ctx)
	if err != nil {
		return nil, 0
	}
	for _, item := range searchResult.Each(reflect.TypeOf(map[string]any{})) {
		rtn = append(rtn, item.(map[string]any))
	}
	return rtn, int(searchResult.TotalHits())
}

func (self *SearchClient) AddFileDocument(data *WendaFile) (err error) {
	ctx := context.Background()

	indexSrv := self.Client.Index().Index(self.Index)
	indexSrv = indexSrv.BodyJson(data)
	res, err := indexSrv.Do(ctx)

	if err != nil {
		// Handle error
		logtt.Print(res, err)
		return  err
	}
	return err
}

// 根据关键key，查询es数据
func (self *SearchClient) GetDocumentByKey(key string, value string) (data []any, err error)  {
	ctx := context.Background()

	matchQuery := elastic.NewMatchQuery(key, value)
	searchResult, err := self.Client.Search().
		Index(self.Index).
		Query(matchQuery).
		Do(ctx)
	if err != nil {
		logtt.Print("ELASTIC.Search.MATCH", err)
	}

	if searchResult == nil {
		return data, err
	}

	var resList []any
	hits := searchResult.Hits.Hits
	for _, val := range hits{
		source := val.Source
		vm := map[string]any{}

		err = json.Unmarshal(source, &vm)
		if err != nil {
			return data, err
		}
		vm["_id"] = val.Id

		resList = append(resList, vm)
	}

	return resList, err
}

func (self *SearchClient) SelectDocumentList(pageInfo request.DatasetFileList) (list []map[string]any, total int64, err error) {
	ctx := context.Background()

	var resList []map[string]any

	var offset, limit int
	limit = pageInfo.PageSize
	if pageInfo.Page == 1 {
		offset = 0
	} else {
		offset = pageInfo.Page * pageInfo.PageSize - 1
	}


	searchResult, err := self.Client.Search().
		Index(self.Index).
		From(offset).Size(limit).
		Do(ctx)
	if err != nil {
		return nil, 0, err
	}
	for _, item := range searchResult.Each(reflect.TypeOf(map[string]any{})) {
		resList = append(resList, item.(map[string]any))
	}

	total = int64(searchResult.TotalHits())
	return resList, total, err
}